from graphviz import Digraph

# 创建一个有向图对象
dot = Digraph(comment='Bert-base-chinese Model', node_attr={'fontname': 'WenQuanYi Zen Hei'}, edge_attr={'fontname': 'WenQuanYi Zen Hei'})
dot.attr(rankdir='TB', size='8,8', dpi='300')

# 定义节点
dot.node('input', '输入层\n（词嵌入 + 位置嵌入 + 句子嵌入）')
dot.node('encoder_block', 'Transformer 编码器块', shape='box3d')
dot.node('output', '输出层\n（根据任务而定）')

# 定义 Transformer 编码器块内的子节点
with dot.subgraph(name='cluster_encoder') as c:
    c.attr(label='Transformer 编码器块', style='filled', color='lightblue')
    c.node('multi_head_attention', '多头注意力机制')
    c.node('feed_forward', '前馈神经网络')
    c.edge('multi_head_attention', 'feed_forward', label='连接')

# 定义边
dot.edge('input', 'encoder_block', label='输入')
dot.edge('encoder_block', 'output', label='输出')

# 渲染图形
dot.render('bert_base_chinese_model.gv', view=True)